import gc
import os
import time
import numpy as np
import psutil
from astropy.io import fits
from astropy.table import Table, join, vstack
from scipy.ndimage import binary_erosion, binary_dilation
from scipy.stats import trim_mean, sigmaclip
from matplotlib import pyplot as plt
from tqdm import trange

from .utils import get_block_limits


def get_cte_data(fname, imgconf, fdark=None):
    data = fits.getdata(fname)
    if fdark is not None:
        data = data - fits.getdata(fdark)

    channels, sdata, snoise, pdata, pnoise, flat, sodata, sonoise, podata, ponoise = [], [], [], [], [], [], [], [], [], []
    maxp = np.max(np.array(imgconf.raw_ny)-np.array(imgconf.y1)-np.array(imgconf.ny))+10
    maxs = np.max(np.array(imgconf.raw_nx)-np.array(imgconf.x1)-np.array(imgconf.nx))+10
    for chan in range(imgconf.nchan):
        x1 = imgconf.raw_x0[chan]
        x2 = x1 + imgconf.raw_nx[chan]
        y1 = imgconf.raw_y0[chan]
        y2 = y1 + imgconf.raw_ny[chan]
        dd = data[y1:y2, x1:x2]
        if np.isnan(dd).sum() > 0:
            continue
        if imgconf.raw_xflip[chan]:
            dd = dd[:, ::-1]
        if imgconf.raw_yflip[chan]:
            dd = dd[::-1, :]
        channels.append(chan+1)
        x1 = imgconf.x1[chan]
        x2 = x1 + imgconf.nx[chan]
        y1 = imgconf.y1[chan]
        y2 = y1 + imgconf.ny[chan]
        ss = dd[y1:y2, x2-10:]
        so = dd[y2:, x2-10:]
        sdata_chan = trim_mean(ss, 0.1, axis=0)
        snoise_chan = np.std(ss , axis=0)
        sodata_chan = trim_mean(so, 0.1, axis=0)
        sonoise_chan = np.std(so, axis=0)
        ########20250207 罗尊修改，为了适应44所器件不同通道过扫区长度不同，保存"过扫区的过扫区数据供进一步分析"
        if len(sdata_chan)<maxs:
            sdata_chan = np.append(sdata_chan,[np.median(sdata_chan[-5:])]*(maxs-len(sdata_chan)))
            snoise_chan = np.append(snoise_chan,[np.median(snoise_chan[-5:])]*(maxs-len(snoise_chan)))
            sodata_chan = np.append(sodata_chan,[np.median(sodata_chan[-5:])]*(maxs-len(sodata_chan)))
            sonoise_chan = np.append(sonoise_chan,[np.median(sonoise_chan[-5:])]*(maxs-len(sonoise_chan)))
        print(len(sdata_chan),len(snoise_chan))
        sdata.append(sdata_chan)
        snoise.append(snoise_chan)
        sodata.append(sodata_chan)
        sonoise.append(sonoise_chan)
        ss = dd[y2-10:, x1:x2]
        so = dd[y2-10:, x2:]
        pdata_chan = trim_mean(ss, 0.1, axis=1)
        pnoise_chan = np.std(ss, axis=1)
        podata_chan = trim_mean(so, 0.1, axis=1)
        ponoise_chan = np.std(so, axis=1)
        if len(pdata_chan)<maxp:
            pdata_chan = np.append(pdata_chan,[np.median(pdata_chan[-5:])]*(maxp-len(pdata_chan)))
            pnoise_chan = np.append(pnoise_chan,[np.median(pnoise_chan[-5:])]*(maxp-len(pnoise_chan)))
            podata_chan = np.append(podata_chan,[np.median(podata_chan[-5:])]*(maxs-len(podata_chan)))
            ponoise_chan = np.append(ponoise_chan,[np.median(ponoise_chan[-5:])]*(maxs-len(ponoise_chan)))
        pdata.append(pdata_chan)
        pnoise.append(pnoise_chan)
        podata.append(podata_chan)
        ponoise.append(ponoise_chan)
        #################################################
        ss = dd[y1:y2, x1:x2]
        flat.append(np.median(ss))

    tab = Table()
    tab['fname'] = [os.path.basename(fname), ] * len(channels)
    tab['channel'] = channels
    tab['serial'] = np.array(sdata).astype('float32')
    tab['serial_sigma'] = np.array(snoise).astype('float32')
    tab['serial_over'] = np.array(sodata).astype('float32')
    tab['serial_over_sigma'] = np.array(sonoise).astype('float32')
    tab['parallel'] = np.array(pdata).astype('float32')
    tab['parallel_sigma'] = np.array(pnoise).astype('float32')
    tab['parallel_over'] = np.array(podata).astype('float32')
    tab['parallel_over_sigma'] = np.array(ponoise).astype('float32')
    tab['flat'] = np.array(flat).astype('float32')
    
    return tab


def plot_cte_data(ftab, ffig):
    xwidth = 100
    ywidth = 200
    tab = Table.read(ftab, format='fits')
    n = len(tab)
    plt.figure(figsize=(8, 4 * n))
    for i in range(n):
        el = tab[i]
        plt.subplot(n, 2, i * 2 + 1)
        plt.plot(el['serial'][10:], color='r')
        plt.plot(el['serial'][10:-xwidth], color='g', zorder=2)
        plt.ylabel('serial overscan (ADU)')
        plt.title('channel {}: {:.1f} ADU'.format(el['channel'], el['flat']))
        plt.subplot(n, 2, i * 2 + 2)
        plt.plot(el['parallel'][10:], color='r')
        plt.plot(el['parallel'][10:-ywidth], color='g', zorder=2)
        plt.ylabel('parallel overscan (ADU)')
        plt.title(os.path.basename(el['fname']))
    plt.tight_layout()
    plt.savefig(ffig)
    plt.close()

def calculate_cte(ftab, fout, imgconf):
    tab = Table.read(ftab)
    flat_serial, flux_serial, fluxerr_serial, n_serial = [], [], [], []
    flat_parallel, flux_parallel, fluxerr_parallel, n_parallel = [], [], [], []
    for el in tab:
        chan = el['channel'] - 1
        xwidth = imgconf.overscan_x2[chan] - imgconf.overscan_x1[chan]
        ywidth = imgconf.overscan_y2[chan] - imgconf.overscan_y1[chan]
        data = el['serial'] - el['serial'][-xwidth:].mean()
        flat_serial.append(np.median(data[:10]))
        flux_serial.append(data[10:30].sum())
        fluxerr_serial.append(np.sqrt((el['serial_sigma'][10:30] ** 2).sum()))
        n_serial.append(imgconf.nx[chan])

        data = el['parallel'] - el['parallel'][-ywidth:].mean()
        flat_parallel.append(np.median(data[:10]))
        flux_parallel.append(data[10:30].sum())
        fluxerr_parallel.append(np.sqrt((el['parallel_sigma'][10:30] ** 2).sum()))
        n_parallel.append(imgconf.ny[chan])

    cat = Table()
    cat['fname'] = tab['fname']
    cat['channel'] = tab['channel']
    cat['texp'] = tab['texp']
    cat['flat'] = tab['flat']

    cat['flux_serial'] = flux_serial
    cat['fluxerr_serial'] = fluxerr_serial
    cat['n_serial'] = n_serial
    cat['flat_serial'] = flat_serial
    cat['cte_serial'] = 1 - cat['flux_serial'] / cat['flat_serial'] / cat['n_serial']
    cat['cte_err_serial'] = cat['fluxerr_serial'] / cat['flat_serial'] / cat['n_serial']
    
    cat['flux_parallel'] = flux_parallel
    cat['fluxerr_parallel'] = fluxerr_parallel
    cat['n_parallel'] = n_parallel
    cat['flat_parallel'] = flat_parallel
    cat['cte_parallel'] = 1 - cat['flux_parallel'] / cat['flat_parallel'] / cat['n_parallel']
    cat['cte_err_parallel'] = cat['fluxerr_parallel'] / cat['flat_parallel'] / cat['n_parallel']
    
    cat.sort(['texp', 'channel'])
    cat.write(fout, format='ipac', overwrite=True)

def group_cte_result(fcat, fgain, fout):
    cat = Table.read(fcat, format='ipac')
    gtab = Table.read(fgain, format='ipac')
    # support both 'gain' and 'gains' column names in gain table
    if ('gain' not in gtab.colnames) and ('gains' in gtab.colnames):
        gtab['gain'] = gtab['gains']
    cc = join(cat, gtab['channel', 'gain'], keys='channel')
    cc['flux'] = cc['flat'] * cc['gain']

    texp = np.unique(cat['texp'])
    flux, cte_s_avg, cte_s_min, cte_p_avg, cte_p_min = [], [], [], [], []
    for t in texp:
        index = cc['texp'] == t
        flux.append(np.mean(cc['flux'][index]))
        cte_s_avg.append(np.mean(cc['cte_serial'][index]))
        cte_s_min.append(np.min(cc['cte_serial'][index]))
        cte_p_avg.append(np.mean(cc['cte_parallel'][index]))
        cte_p_min.append(np.min(cc['cte_parallel'][index]))

    tab = Table()
    tab['texp'] = texp
    tab['flux'] = np.array(flux).astype('float32')
    tab['flux'].unit = 'ph'
    tab['serial_avg'] = np.array(cte_s_avg) * 100
    tab['serial_avg'].unit = '%'
    tab['serial_min'] = np.array(cte_s_min) * 100
    tab['serial_min'].unit = '%'
    tab['parallel_avg'] = np.array(cte_p_avg) * 100
    tab['parallel_avg'].unit = '%'
    tab['parallel_min'] = np.array(cte_p_min) * 100
    tab['parallel_min'].unit = '%'
    tab.write(fout, format='ipac', overwrite=True)

def _fit_baseline(data):
    x = np.arange(len(data))
    sigma = sigmaclip(data, 5, 5)[0].std()
    thresh = np.median(data) + 5 * sigma
    index = data > thresh
    index = binary_erosion(index, iterations=3)
    valid = ~binary_dilation(index, iterations=6)
    coef = np.polyfit(x[valid], data[valid], deg=2)
    y = np.poly1d(coef)(x)
    sigma = sigmaclip(data - y, 5, 5)[0].std()
    index = binary_erosion((data - y) < sigma * 5, iterations=1) * valid
    coef = np.polyfit(x[index], data[index], deg=3)
    return coef


def get_trap_image(fname, fgain, fout, imgconf, fdark=None):
    hdulist = fits.HDUList([fits.PrimaryHDU(), ])
    hdulist[0].header['FNAME'] = os.path.basename(fname)

    tab = Table.read(fgain, format='ipac')
    # support both 'gain' and 'gains' column names in gain table
    if ('gain' not in tab.colnames) and ('gains' in tab.colnames):
        tab['gain'] = tab['gains']
    with fits.open(fname) as img:
        if fdark is not None:
            data = img[0].data - fits.getdata(fdark)
        else:
            data = img[0].data

        for chan in range(imgconf.nchan):
            if not ((chan + 1) in tab['channel']):
                continue
            index = tab['channel'] == chan + 1
            gain = tab['gain'][index][0]

            x1 = imgconf.raw_x0[chan]
            x2 = x1 + imgconf.raw_nx[chan]
            y1 = imgconf.raw_y0[chan]
            y2 = y1 + imgconf.raw_ny[chan]
            dd = data[y1:y2, x1:x2]
            if imgconf.raw_xflip[chan]:
                dd = dd[:, ::-1]
            if imgconf.raw_yflip[chan]:
                dd = dd[::-1, :]

            x1 = imgconf.x1[chan]
            x2 = x1 + imgconf.nx[chan]
            y1 = imgconf.y1[chan]
            y2 = y1 + imgconf.ny[chan]
            ss = dd[y2:, x1:x2] * gain

            hdu = fits.ImageHDU(ss)
            hdu.header['FNAME'] = os.path.basename(fname)
            hdu.header['CHANNEL'] = chan + 1
            hdu.header['GAIN'] = gain
            hdu.header['FLUX'] = np.median(dd[y1:y2, x1:x2]) * gain
            hdu.name = 'RAW{}'.format(chan+1)
            hdulist.append(hdu)

            ny, nx = ss.shape
            dd = ss.mean(axis=0)
            coef = _fit_baseline(dd)
            model = np.poly1d(coef)
            bkg = np.transpose(np.repeat(model(np.arange(nx)), ny).reshape([nx, ny]))
            hdu = fits.ImageHDU(ss - bkg, header=hdulist[-1].header)
            hdu.header['POLY0'] = coef[0]
            hdu.header['POLY1'] = coef[1]
            hdu.header['POLY2'] = coef[2]
            hdu.header['POLY3'] = coef[3]
            hdu.name = 'RED{}'.format(chan+1)
            hdulist.append(hdu)

    hdulist.writeto(fout, overwrite=True)


def bin_trap_data(data):
    dd = data.sum(axis=0)
    dd2 = data[0, :]
    return dd, dd2


def get_trap_position(fdata, fout, ffig=None):
    with fits.open(fdata) as img:
        n = (len(img) - 1) // 2

        channels, flats, gains, pos, value, value_first = [], [], [], [], [], []
        channels.append(-999)
        flats.append(np.float32(-999))
        gains.append(np.float32(-999))
        pos.append(-999)
        value.append(np.float32(-999))
        value_first.append(np.float32(-999))

        plt.figure(figsize=(10, 4 * n))

        for i in range(n):
            header = img[i * 2 + 2].header
            print("**********img[i * 2 + 2].data.shape",img[i * 2 + 2].data.shape)
#            d1 = img[i * 2 + 2].data.sum(axis=0)    #forgot to choose first 10 rows
            d1 = img[i * 2 + 2].data[:10,:].sum(axis=0)
            print("**********d1.shape",d1.shape)
            d2 = img[i * 2 + 2].data[0, :]
            x = np.arange(len(d1))
            #print("x",x)
            index = (d1 > 200) * (d2 > 10)
            index = x[index]
            for p in index:
                channels.append(header['CHANNEL'])
                flats.append(header['FLUX'])
                gains.append(header['GAIN'])
                pos.append(p + 1)  # convert from python to fits standard
                value.append(np.float32(d1[p]))
                value_first.append(np.float32(d2[p]))

            plt.subplot(n, 2, i * 2 + 1)
            coef = [header['POLY0'], header['POLY1'], header['POLY2'], header['POLY3']]
            model = np.poly1d(coef)
            y = model(x) * img[i * 2 + 2].data.shape[0]
            y -= np.median(y)
            plt.plot(x, y, color='coral', zorder=2)
            ax = plt.gca()
            ylim = ax.get_ylim()
            sigma = sigmaclip(d1, 5, 5)[0].std()
            ylim = (ylim[0] - sigma * 3, max(ylim[1] + sigma * 7, 220))
            plt.plot(x, y + d1, color='g', zorder=1)
            for p in index:
                plt.plot([p, p], [ylim[0], ylim[1]], color='r', ls='--', zorder=3, alpha=0.5)
            if len(index) > 0:
                plt.scatter(x[index], (y + d2)[index], color='r', zorder=3)
            plt.grid(True)
            plt.title('channel: {}'.format(header['CHANNEL']))
            plt.ylabel('sum of 10 rows (e-)')
            plt.ylim(ylim)

            plt.subplot(n, 2, i * 2 + 2)
            y = model(x)
            y -= np.median(y)
            plt.plot(x, y, color='coral', zorder=2)
            ax = plt.gca()
            ylim = ax.get_ylim()
            sigma = sigmaclip(d2, 5, 5)[0].std()
            ylim = (ylim[0] - sigma * 3, max(ylim[1] + sigma * 7, 11))
            plt.plot(x, y + d2, color='g', zorder=1)
            for p in index:
                plt.plot([p, p], [ylim[0], ylim[1]], color='r', ls='--', zorder=3, alpha=0.5)
            if len(index) > 0:
                plt.scatter(x[index], (y + d2)[index], color='r', zorder=3)
            plt.grid(True)
            plt.title('flat = {:.1f} e-'.format(flats[-1]))
            plt.ylabel('first row (e-)')
            plt.ylim(ylim)

        plt.tight_layout()
        if ffig is not None:
            plt.savefig(ffig)
        plt.close()

        tab = Table()
        tab['name'] = [os.path.basename(fdata), ] * len(flats)
        tab['channel'] = channels
        tab['gain'] = gains
        tab['gain'].unit = 'ph/adu'
        tab['flat'] = flats
        tab['flat'].unit = 'ph'
        tab['position'] = pos
        tab['sum'] = value
        tab['sum'].unit = 'ph'
        tab['first'] = value_first
        tab['first'].unit = 'ph'
        tab.sort(['channel', 'position'])
        tab.write(fout, format='ipac', overwrite=True)


def merge_trap_cat(flist, fout):
    cat = Table()
    for f in flist:
        tab = Table.read(f, format='ipac')
        cat = vstack([cat, tab])
    index = cat['channel'] > 0
    cat = cat[index]
    if len(cat) > 0:
        fnames = np.unique(cat['name'])
        flats = []
        for f in fnames:
            index = cat['name'] == f
            flats.append(np.mean(cat['flat'][index]))
        index = np.argmin(np.abs(np.array(flats) - 1000))
        fname = fnames[index]
        index = cat['name'] == fname
        cat = cat[index]
    cat.write(fout, format='ipac', overwrite=True)


def get_cmos_trap(flist, fbias, fgain, fout, fflat=None, scale=None, malloc=None):
    # memory management
    if malloc is None:
        malloc = psutil.virtual_memory().available / 1024 ** 3 / 2

    # initial data
    n = len(flist)
    if n == 0:
        raise Exception('No valid data for trap mask.')
    bias = fits.getdata(fbias).astype(np.float32)
    gain = fits.getdata(fgain).astype(np.float32)
    ny, nx = bias.shape
    output = np.zeros((ny, nx), dtype=np.int16)
    output_flat = None
    if fflat is not None:
        output_flat = np.zeros((ny, nx), dtype=np.float32)
    if scale is None:
        scale = np.ones(n).astype(np.float32)

    # decide which rows to process
    ncols = malloc / (nx * 8 / 1024 ** 3) / 2 / n
    if ncols < 2:
        raise Exception('Error: low memory.')
    nblocks = int(np.floor(ny / ncols)) + 1
    ylims = get_block_limits(ny, nblocks)

    for i in range(nblocks):
        y1, y2 = ylims[i], ylims[i + 1]
        cutout = np.zeros((y2 - y1, nx, n), dtype=np.float32)
        bb = bias[y1:y2, :]
        gg = gain[y1:y2, :]
        t = time.asctime(time.localtime(time.time()))
        for j in trange(n, desc='{}: block {}/{}'.format(t, i+1, nblocks)):
            with fits.open(flist[j]) as img:
                cutout[:, :, j] = (img[0].data[y1:y2, :] - bb) * gg / scale[j]
        mmn = np.mean(cutout, axis=2)
        if fflat is not None:
            output_flat[y1:y2, :] = mmn.astype(np.float32)
        mmn = np.repeat(mmn, n).reshape(cutout.shape)
        index = np.abs(cutout - mmn) > 200
        index = index.sum(axis=2) > 1
        output[y1:y2, :] = index.astype(np.int16) * 128
        del cutout, bb, gg
        gc.collect()

    hdu = fits.PrimaryHDU(output)
    hdu.writeto(fout, overwrite=True)
    if fflat is not None:
        hdu = fits.PrimaryHDU(output_flat)
        hdu.writeto(fflat, overwrite=True)
